Add Megatron-Bridge prune & quantize launcher pipelines - #2031
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds post-run accuracy and pruning-score gates. It also adds inline launcher commands, per-task requirements, configurable Docker users, cache environment overrides, workflow examples, documentation, and validation tests. ChangesEvaluation and pruning gates
Launcher task workflows
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant EvaluationCLI
participant EvalRunner
participant ResultsJSON
EvaluationCLI->>EvaluationCLI: validate threshold, output path, and one task
EvaluationCLI->>EvalRunner: execute evaluation
EvalRunner->>ResultsJSON: write results JSON
EvaluationCLI->>ResultsJSON: read task accuracy
ResultsJSON-->>EvaluationCLI: return accuracy
EvaluationCLI-->>EvaluationCLI: report PASS or exit status 1
sequenceDiagram
participant SandboxPipeline
participant DockerExecutor
participant TaskShell
SandboxPipeline->>SandboxPipeline: resolve task fields
SandboxPipeline->>DockerExecutor: configure Docker user
DockerExecutor->>TaskShell: run task
TaskShell->>TaskShell: install requirements
TaskShell-->>DockerExecutor: execute inline or script command
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
/claude review |
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
SCOPE: Reviewed all 9 files in the PR authoritative changed-file list: tools/launcher/{core.py, slurm_config.py, docs/configuration.md}, the two new example YAMLs, three test files, and examples/llm_eval/lm_eval_hf.py. Pure tooling/examples change — no modelopt/ source, mode registration, config schema, or export code touched, so ModeState/Export/Algorithm categories do not apply. (A raw two-dot diff against the shallow base also surfaced unrelated MiniMax deletions and hf_ptq README edits; those are NOT in the PR file list and were treated as base-mismatch noise, not reviewed.)
FINDINGS: CRITICAL 0, IMPORTANT 0, SUGGESTION 1.
- SUGGESTION: lm_eval_hf.py accuracy gate uses sorted(glob(...))[-1] to pick the newest results file, which orders by path (directory name first), not timestamp. Robust only when --output_path is not reused across runs; shipped examples use ephemeral per-run paths so non-blocking. Suggested max(files, key=os.path.getmtime).
ASSESSMENT: Low risk. Two substantive paths check out:
- core.py inline/reqs/docker_user — new SandboxTask fields (inline, reqs, reqs_file) and SlurmConfig.docker_user default to None, slurm_factory gains an optional kwarg — backward compatible. shlex imported; reqs tokens shlex-quoted so specifiers like transformers<5 stay literal; args+inline correctly rejected; docker_user falls back to host uid:gid.
- lm_eval_hf.py gate + backend handling — for non-hf backends ModelOpt-only kwargs are dropped from the namespace (not injected into model_args), so --model vllm on a quantized checkpoint will not break the constructor. Metric extraction matches acc / acc, while excluding acc_stderr / acc_norm; single-task validation fails fast before the expensive eval.
Tests cover the new fields, the guard, and docker_user override. LGTM.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2031 +/- ##
===========================================
+ Coverage 66.94% 77.79% +10.84%
===========================================
Files 519 520 +1
Lines 59401 60599 +1198
===========================================
+ Hits 39767 47144 +7377
+ Misses 19634 13455 -6179
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
ac86810 to
e8b8914
Compare
|
6c137be to
9d4ee18
Compare
9d4ee18 to
4aef972
Compare
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 8
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/llm_eval/lm_eval_hf.py`:
- Around line 339-351: Update the lower_bound validation branch around
_pop_gate_bound and _enforce_accuracy_gate to require a non-empty output_path
before cli.execute(args) runs. Raise the existing appropriate validation error
immediately when --accuracy_lower_bound is used without --output_path, while
preserving the single-task validation and normal execution path.
In `@tests/examples/megatron_bridge/test_prune_minitron.py`:
- Line 73: Strengthen the score-gate coverage in the real-script tests around
the existing score_lower_bound cases by adding an invocation with an impossible
MMLU bound such as 1.1. Assert that this command exits non-zero, while
preserving the existing successful-path checks for valid bounds.
In `@tools/launcher/core.py`:
- Around line 788-791: Enforce the inline-task command-source contract before
command construction in tools/launcher/core.py around the existing task.inline
validation: reject tasks that define both script and inline, and reject tasks
that define neither, while preserving the existing inline args restriction.
Update tools/launcher/tests/test_examples_resolve.py in the example validation
logic to reject args when inline is configured, matching runtime behavior.
- Around line 848-852: Update the marker construction in the launcher command
that defines reqs_prefix to remove the # nosec suppression and derive a task-
and node-unique synchronization path using the available job/task and node
identifiers. Ensure the resulting path remains safely rooted in the temporary
directory and is shared by the intended local ranks, while preserving the
existing wait-and-touch synchronization behavior.
- Around line 855-857: Update the reqs_prefix branch around script_cmd
construction to shell-quote task.script and every value in task_args before
concatenating them into the inline bash command. Preserve the existing argument
order and run.Script invocation while ensuring YAML/global-variable contents are
passed literally rather than interpreted by the shell.
In `@tools/launcher/docs/configuration.md`:
- Around line 63-66: Update the inline task documentation to state that args
must be omitted when inline is used, because run_jobs() raises ValueError if
inline is combined with non-empty args; remove the claim that args are ignored
while preserving the existing inline command behavior description.
In
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml`:
- Around line 19-20: Add the required environment variables to both model
configurations: in
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml
lines 19-20, add MLM_MODEL_CFG using the Nemotron HuggingFace repository ID; in
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml
lines 16-17, add MLM_MODEL_CFG with the same repository ID and QUANT_CFG
matching the CLI quantization configuration.
In `@tools/launcher/tests/test_examples_resolve.py`:
- Around line 79-91: Update the task-shape validation in the test around the
script/inline exclusivity assertions to reject any task that combines a
non-empty inline command with args, matching run_jobs() behavior. Add an
assertion covering inline tasks’ args absence so the test exercises this
runtime-invalid shape while preserving existing script/inline and single-line
checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4b2432c3-8b73-47b4-85a4-b0f1bbdc814f
📒 Files selected for processing (12)
examples/llm_eval/lm_eval_hf.pyexamples/megatron_bridge/prune_minitron.pymodelopt/torch/prune/plugins/mcore_minitron.pytests/examples/megatron_bridge/test_prune_minitron.pytools/launcher/core.pytools/launcher/docs/configuration.mdtools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yamltools/launcher/slurm_config.pytools/launcher/tests/test_docker_execution.pytools/launcher/tests/test_examples_resolve.pytools/launcher/tests/test_yaml_formats.py
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Mostly a tidy set of additive launcher/example features, and the E2E evidence in the PR body is appreciated. A few substantive issues, though:
Correctness / backward compat
- Adding
"best"toMCoreMinitronSearcher.default_state_dictis not backward compatible for cached score checkpoints:BaseSearcher.load_search_checkpoint()is called withstrict=Trueand assertsstate_dict.keys() == self.state_dict().keys(), so any pre-existing--prune_intermediate_ckpt/modelopt_pruning_scoresdir written by an older version now fails withKeys in checkpoint don't match!instead of resuming. The PR body claims "new state-dict key is additive" — it is additive for writers, not readers. self.bestshadowsBaseSearcher.best, which is the documented search-result state dict returned byBaseSearcher.search() -> SearchStateDict(reset_searchsetsself.best: SearchStateDict = {}).MCoreMinitronSearcher.search()now returns aCandidateSubnet | None, breaking that annotation/contract (and likely tripping mypy on the attribute override). Please rename (e.g.best_candidate).lm_eval_hf.pynow silently drops--quant_cfg/--sparse_cfg/--auto_quantize_bitsfor non-hfbackends. A user running--model vllm --quant_cfg FP8_DEFAULT_CFGgets an eval of the unquantized model with no warning — and with--accuracy_lower_boundthat even reports PASS. This should raise instead of silently ignoring. Alsois_hfmisses thehuggingfacealias andhf-multimodal, both of which are HFLM-based.
Tests
4. The new gate helpers (_requested_tasks, _enforce_accuracy_gate) are pure functions with real edge cases (mtime selection among nested results*.json, "acc,none" vs acc_stderr, missing task) and have no tests. Similarly, the new run_jobs logic (reqs_prefix shell string, args+inline ValueError) is untested — the added test_examples_resolve assertion covers YAML-level script/inline exclusivity, not the args+inline guard the PR body claims.
5. get_default_env now reads ambient HF_HOME/TRITON_CACHE_DIR, but the existing tests/test_core.py::TestGetDefaultEnv assertions ("/cicd/hf-cache") weren't updated to monkeypatch.delenv, so they'll fail on any dev box/CI runner that exports those vars. No test covers the new override path either.
Robustness / docs
6. The pip barrier marker /tmp/.modelopt_launcher_reqs_done is not run-unique, and under enroot/pyxis /tmp is commonly the host /tmp; a stale marker from an earlier task on the same node makes non-zero local ranks skip the wait, reintroducing exactly the concurrent-pip race the barrier guards against. Ranks also proceed after the 600 s timeout even when rank 0's install failed.
7. docs/configuration.md says args "is ignored" with inline, but run_jobs raises ValueError.
Design (per the complexity gate): inline/reqs extend the existing launcher config rather than adding a new subsystem, which is good, but they do add a third way to express a task command alongside the documented common/**/*.sh wrapper convention (CLAUDE.md: "Scripts go in common/") and typed _target_ tasks — and common/megatron_bridge/import/import.sh shows the wrapper pattern already covers this flow (a wrapper could also do the pip install + rank-0 barrier in bash, making reqs/reqs_file unnecessary). The PR body describes the new fields but doesn't say why the wrapper-script or typed-task route was rejected; please add that rationale, and update CLAUDE.md's conventions list so contributors know which of the three to reach for.
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/megatron_bridge/prune_minitron.py`:
- Around line 683-693: The score-gate logic currently runs after checkpoint
export, allowing failed pruned models to remain at the output path and bypass
evaluation on reruns. Move the `args.score_lower_bound` validation block,
including the `pruning_scores["best"]` check and `sys.exit(1)`, to immediately
after `mtp.prune()` and before either artifact-save path.
In `@tools/launcher/core.py`:
- Around line 849-852: Update the reqs_prefix shell flow around marker creation
so nonzero local ranks fail when the marker is absent after the wait loop,
preventing fi && task_cmd from running without installed dependencies. Preserve
successful continuation when {marker} appears and propagate an explicit failure
status when the 600-iteration wait expires.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2e1eb71d-dedb-4e11-adee-27243a5a5b6e
📒 Files selected for processing (12)
examples/llm_eval/lm_eval_hf.pyexamples/megatron_bridge/prune_minitron.pymodelopt/torch/prune/plugins/mcore_minitron.pytests/examples/megatron_bridge/test_prune_minitron.pytools/launcher/core.pytools/launcher/docs/configuration.mdtools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yamltools/launcher/slurm_config.pytools/launcher/tests/test_docker_execution.pytools/launcher/tests/test_examples_resolve.pytools/launcher/tests/test_yaml_formats.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tools/launcher/docs/configuration.md
4aef972 to
d4cf77f
Compare
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 6
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/llm_eval/lm_eval_hf.py`:
- Line 267: Update the argparse definitions for calib_size,
auto_quantize_method, and auto_quantize_score_size to default to None, so non-HF
backends only reject values explicitly supplied by the user. In the HF
model_args injection path, apply the existing effective defaults (512,
"gradient", and 128) when these options remain unset, while preserving the
current validation for explicitly provided values.
In `@modelopt/torch/prune/plugins/mcore_minitron.py`:
- Line 645: Update the assignment in the searcher implementation from self.best
to a distinct field such as self.best_candidate when storing asdict(best).
Preserve BaseSearcher.self.best exclusively for its SearchStateDict so the
inherited search() contract remains unchanged.
- Line 645: Update the search completion flow around self.best = asdict(best) to
save the search checkpoint after assigning the selected candidate, ensuring the
final checkpoint’s "best" state is populated. Use the existing
save_search_checkpoint() mechanism and preserve the earlier checkpoint behavior.
- Line 297: The new "best" key added to the checkpoint structure causes strict
validation to fail when loading older checkpoints that lack this field. Either
exclude the "best" field from the strict state validation in
BaseSearcher.load_search_checkpoint, or add migration logic that populates the
"best" key in older checkpoints before strict validation occurs. Ensure backward
compatibility so existing checkpoints can resume without modification.
In
`@tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yaml`:
- Around line 20-34: Remove the hardcoded --trust_remote_code flags from both
quantization and export commands in the example; remote-code execution must
default to False and require explicit caller opt-in, with the documented
security-exception approval obtained if it remains necessary.
- Around line 17-24: Update the environment section to define MLM_MODEL_CFG as
the Hugging Face repository ID and QUANT_CFG as MAMBA_MOE_FP8_CONSERVATIVE_CFG,
then replace the hardcoded model path and quantization config in the inline
quantize.py command with these variables.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d767c065-f193-45ba-8fca-60d5f62bb30f
📒 Files selected for processing (13)
examples/llm_eval/lm_eval_hf.pyexamples/megatron_bridge/prune_minitron.pymodelopt/torch/prune/plugins/mcore_minitron.pytests/examples/megatron_bridge/test_prune_minitron.pytools/launcher/core.pytools/launcher/docs/configuration.mdtools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yamltools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_quantize.yamltools/launcher/slurm_config.pytools/launcher/tests/test_core.pytools/launcher/tests/test_docker_execution.pytools/launcher/tests/test_examples_resolve.pytools/launcher/tests/test_yaml_formats.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tests/examples/megatron_bridge/test_prune_minitron.py
- examples/megatron_bridge/prune_minitron.py
- tools/launcher/tests/test_examples_resolve.py
- tools/launcher/slurm_config.py
- tools/launcher/tests/test_yaml_formats.py
- tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/mbridge_prune.yaml
- tools/launcher/docs/configuration.md
End-to-end ModelOpt launcher pipelines for Megatron-Bridge on Nemotron-3-Nano-30B-A3B: - mbridge_prune.yaml: prune -> vLLM gen -> MMLU gate - mbridge_quantize.yaml: FP8 quantize -> unified-HF export -> vLLM gen -> MMLU gate (vLLM backend) Launcher (tools/launcher) additions so these run wrapper-free from YAML: - SandboxTask.inline: run a command directly from YAML (no .sh wrapper) - SandboxTask.reqs / reqs_file: pip-install deps in-container before the command - SlurmConfig.docker_user: run local Docker as a chosen user (e.g. root) - reject args together with inline examples/llm_eval/lm_eval_hf.py: - --accuracy_lower_bound: gate the run on a single task's acc (exit non-zero if below) - drop ModelOpt args for non-hf backends so --model vllm works on quantized checkpoints Docs (tools/launcher/docs/configuration.md) + unit tests updated. Verified end-to-end on Qwen3-0.6B (prune, FP8 quantize/export, vLLM gen, MMLU gates). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
d4cf77f to
b090a7c
Compare
|
Addressed the first round of review comments and resolved those threads (changes in Prune score gate —
Not addressed yet (will follow up): the test-coverage suggestions (failing-gate assertion, |
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/examples/llm_eval/test_llm_eval.py`:
- Around line 39-41: Add a focused regression test alongside the existing LLM
evaluation accuracy test that configures an accuracy lower bound above the
produced score, invokes the relevant lm_eval_hf.py execution path, and asserts a
non-zero result. Keep the existing passing-bound test unchanged and ensure the
new test specifically exercises the failing accuracy-gate behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d9a9085f-2967-45cb-975a-9cda12259235
📒 Files selected for processing (1)
tests/examples/llm_eval/test_llm_eval.py
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Most of the previous round is genuinely resolved (see below); what's left are owner-judgment calls plus one test gap, so I'd like a human sign-off rather than auto-approving.
Confirmed fixed: self.best = asdict(best) preserves the BaseSearcher.search() -> SearchStateDict contract (base returns self.best; reset_search()'s comment anticipates best in default_state_dict; mtn.search() returns state_dict(), so pruning_scores["best"]["score"] resolves); non-HF backends now raise on quantization/sparsity request args and is_hf covers all four HFLM aliases; results file selected by mtime; --accuracy_lower_bound validates --output_path/single-task before the expensive eval; TestGetDefaultEnv env isolation added; docs say args is rejected (not ignored); the reqs barrier now rm -fs the marker and waiting ranks fail via [ -f {marker} ]; script/inline exclusivity enforced in run_jobs.
Still warranting a look:
-
💬 Author decided to accept the strict-checkpoint-key break (
"best"added todefault_state_dict) because pruning is a one-shot step — reasonable, but flagging for human sign-off because the failure mode for anyone resuming a pre-upgrade--prune_intermediate_ckpt/modelopt_pruning_scoresdir is an opaqueassert state_dict.keys() == self.state_dict().keys(), "Keys in checkpoint don't match!"with no hint to delete the dir, and the mitigation is one line (overrideload_search_checkpoint(strict=False)inMCoreMinitronSearcher, or derive the score from the already-presentall_candidates_per_constraint). Owner's call, but worth a conscious yes/no. -
💬 Author declined new tests for the
run_jobsguard as "a one-line guard" — the guard itself, agreed, but thereqs_prefixconstruction (shlex quoting ofreqs/reqs_file, the rank-0 barrier shell string, and thereqs+script→run.Script(inline=...)branch) is the most intricate new code in this PR and is currently exercised by nothing.tests/test_core_extended.py::TestRunJobsExtendedalready patchescore.run.Scriptand asserts on the call kwargs, so a test asserting theinline=string for areqs/inlinetask is ~10 lines. Both failing-gate paths (--accuracy_lower_boundabove the score,--score_lower_boundabove the score) also remain untested; the prune one is fairly argued as too expensive, but thelm_evalone could be a cheap unit test of_enforce_accuracy_gateagainst a syntheticresults.json. -
💬 Design rationale (raised previously): the PR body describes
inline/reqsbut still doesn't say why the existingcommon/**/*.shwrapper convention (whichcommon/megatron_bridge/import/import.shalready demonstrates for this flow, and which could do the pip install + barrier in bash) or typed_target_tasks were rejected.tools/launcher/CLAUDE.mdstill lists only "Scripts go incommon/" under key conventions, so contributors now have three ways to express a task command with no guidance on which to pick. Please add a sentence of rationale to the PR body and a line toCLAUDE.md(docs/configuration.mdis updated nicely already). -
Multi-node
reqsbarrier (nit): the marker is created in the per-task working dir (/nemo_run/code), which is the packaged code dir mounted from the shared job dir — i.e. shared across nodes, not node-local. On a multi-node task, node B's non-zero local ranks can observe node A's marker and skip the wait before node B's local rank 0 has installed anything. Folding${SLURM_NODEID:-0}into the marker name (or using a node-local path) would make it truly per-node; alternatively documentreqsas single-node only. -
score_lower_bound=0.01intest_prune_minitron(nit): couples the GPU test's pass/fail to a tiny random-init model's MMLU score. Low risk (4-choice MMLU floor is ~25%), but it is a new flake surface for a test whose purpose is the prune path, not accuracy.
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review at head (post-d669c23). Most of the previous round is genuinely fixed; what remains is (a) the design/docs ask that has now been raised twice without a change, and (b) the lm_eval_hf.py gate, which the operator asked me to look at for simplification — it can be materially shorter.
Confirmed resolved since the last round: non-HF backends now raise on quantization/sparsity request args (and only on those — calib_size/auto_quantize_method defaults no longer false-positive); is_hf covers hf/hf-auto/huggingface/hf-multimodal; results file selected by mtime; --accuracy_lower_bound validates --output_path + single-task before cli.execute; script/inline exclusivity + args+inline rejected in run_jobs; script path shell-quoted in the reqs+script branch; reqs barrier rm -fs the marker and waiting ranks fail via [ -f {marker} ]; TestGetDefaultEnv env isolation added; docs/configuration.md documents args as rejected (not ignored) and the new inline/reqs fields nicely.
Answering the specific question — yes, examples/llm_eval/lm_eval_hf.py can be simplified. The gate is currently ~70 added lines across four helpers; I count three reductions that don't lose behavior (details inline):
- The whole
glob+mtime+jsonre-read (and the hard--output_pathrequirement, and its duplicated validation) exists only because the results are read back off disk. If lm-eval'sHarnessCLI.execute(args)returns the results dict (ascli_evaluatehistorically did),results["results"][task]gives the same number in-process and_enforce_accuracy_gatecollapses to ~6 lines with no--output_pathcoupling. Worth 5 minutes to check against the pinned lm-eval ≥ 0.4.10. _pop_gate_boundis a 5-line function forvars(args).pop("accuracy_lower_bound", None);_requested_tasksis a triple-nested comprehension for what is already a normalized--tasksvalue.- The non-HF validation is interleaved with the injection loop and needs a second key set (
quant_request_keys) plus a 4-line comment to explain itself; hoisting it above the loop makes both halves obvious and reports all offending flags at once.
Separately, worth deciding whether this gate belongs in the shared example CLI at all: it only serves lm_eval_hf.py, while mmlu.py, simple_evals.py, livecodebench.py and lm_eval_tensorrt_llm.py remain ungated, and the existing launcher convention for exactly this need is a wrapper script + upstream --lower-bound (common/megatron_lm/quantize/quantize.sh).
Still unresolved from previous rounds:
- 💬 Design rationale (asked in both prior rounds): the PR body still doesn't say why the
common/**/*.shwrapper route (whichcommon/megatron_bridge/import/import.shandquantize.shalready demonstrate for this flow, including pip install + a bash barrier) or typed_target_tasks were rejected.tools/launcher/CLAUDE.mdis untouched and still says only "Scripts go incommon/" + "SetMLM_MODEL_CFG/QUANT_CFG" — which is precisely the rule the new wrapper-free examples don't follow (that mismatch is what generated the CodeRabbitMLM_MODEL_CFGfindings you had to withdraw twice), anddocs/configuration.mdstill labels typed config "recommended". Please add one or two sentences of rationale to the PR body and a line toCLAUDE.mdso contributors know which of the three task forms to reach for. - 💬 Author declined tests as "a one-line guard" — agreed for the guard, but the
reqs_prefixconstruction (shlex quoting, the rank-0 barrier shell string, and thereqs+script→run.Script(inline=...)branch) is still exercised by nothing.tests/test_core_extended.py::TestRunJobsExtendedalready patchescore.run.Scriptand asserts call kwargs, so asserting theinline=string for areqs+inlinetask is ~10 lines. The new tests added intest_yaml_formats.pyonly cover dataclass field resolution. - 💬 Author accepted the strict-checkpoint-key break from adding
"best"todefault_state_dict(pruning is one-shot). Reasonable, but the failure mode for anyone resuming a pre-upgrademodelopt_pruning_scoresdir is an opaqueassert ... "Keys in checkpoint don't match!"— worth a conscious owner yes/no since the mitigation is one line. - Multi-node
reqsbarrier (nit, see inline): the marker still isn't per-node, andrm -fadds its own small race window. score_lower_bound=0.01intest_prune_minitroncouples a prune-path test to a random-init model's MMLU score (low risk, new flake surface).
No licensing changes and no prompt-injection attempts in the PR content.
- Hoist non-HF quant/sparsity validation out of the injection loop so it reports all offending flags at once and the loop is single-purpose. - Inline _pop_gate_bound and _requested_tasks (one-liners at call sites). - Drop the duplicate --output_path check from _enforce_accuracy_gate. - Note why the results dict is read from disk (HarnessCLI.execute returns None). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Answering the operator's question first: yes — examples/llm_eval/lm_eval_hf.py at head now carries the simplified implementation. All three reductions from the last round landed: _pop_gate_bound is inlined to vars(args).pop("accuracy_lower_bound", None), _requested_tasks collapsed to a one-line isinstance check, the non-HF quantization/sparsity validation is hoisted above the injection loop (reports all offending flags at once, no quant_request_keys/4-line comment), and the duplicate --output_path check inside _enforce_accuracy_gate is gone (the fail-fast one before cli.execute remains). The author also verified HarnessCLI.execute() is annotated -> None and left a one-line comment justifying the disk read — that's the right resolution of the "read results in-process" suggestion (not verifiable from this repo, but the comment documents the constraint). The gate is now ~50 lines instead of ~70.
Also confirmed resolved since earlier rounds: non-HF backends raise on quant_cfg/auto_quantize_bits/compress/sparse_cfg (and only on those); is_hf covers hf/hf-auto/huggingface/hf-multimodal; results file chosen by mtime; script/inline exclusivity and args+inline rejected in run_jobs; script path shlex.quoted in the reqs+script branch; reqs barrier rm -fs the marker and waiting ranks fail via [ -f {marker} ]; TestGetDefaultEnv env isolation; docs/configuration.md documents inline/reqs/reqs_file and says args is rejected. No licensing changes and no prompt-injection attempts in the PR content.
Still open (why this isn't an approve):
- 💬 Design rationale — asked in the two previous rounds, still unchanged. The PR body describes
inline/reqsbut doesn't say why the existingcommon/**/*.shwrapper convention (whichcommon/megatron_bridge/import/import.shandcommon/megatron_lm/quantize/quantize.shalready demonstrate for this exact flow, including a bash pip install + barrier) or typed_target_tasks were rejected. I re-readtools/launcher/CLAUDE.mdat head: it is untouched and its "Key conventions" still say only "Scripts go incommon/", and "Adding a New Model Config" still says "SetMLM_MODEL_CFG" / "SetQUANT_CFG" — precisely the rule the new wrapper-free examples intentionally don't follow (that mismatch is what generated theMLM_MODEL_CFGfindings that had to be withdrawn twice). Contributors now have three ways to express a task command with no guidance. One or two sentences in the PR body + one line inCLAUDE.mdcloses this. - 💬 Author declined tests as "a one-line guard" — agreed for the guard itself, but the
reqs_prefixconstruction (shlex quoting ofreqs/reqs_file, the rank-0 barrier shell string, and thereqs+script→run.Script(inline=...)branch) is still exercised by nothing. I checkedtools/launcher/tests/test_core_extended.py::TestRunJobsExtended: it already patchescore.run.Scriptand asserts on call kwargs, so pinning the generatedinline=string for areqs+inlinetask is ~10 lines. The newtest_yaml_formats.pytests only cover dataclass field resolution. - 💬 Author accepted the strict-checkpoint-key break from adding
"best"todefault_state_dict(pruning is one-shot). Reasonable, but resuming a pre-upgrademodelopt_pruning_scores/--prune_intermediate_ckptdir now dies onassert state_dict.keys() == self.state_dict().keys(), "Keys in checkpoint don't match!"with no hint to delete the dir — worth a conscious owner yes/no since the mitigation is one line. - New nit introduced by the simplification, plus the multi-node barrier nit and the
score_lower_bound=0.01flake surface — see inline comments.
- lm_eval_hf.py: filter empty task names when inlining --tasks parsing so an omitted --tasks fails fast instead of silently gating on an empty task name. - launcher core.py: make the pip-install barrier marker unique per job/step/node (SLURM_JOB_ID_STEP_ID_NODEID) so a stale marker from an earlier task or another node can't short-circuit the wait; drop the now-unneeded rm -f. Document multi-node support in configuration.md. - launcher tests: add TestRunJobsExtended cases pinning the generated reqs barrier inline string for both inline and script tasks (quoting + marker shape). - test_prune_minitron.py: use score_lower_bound=0.0 to exercise the gate path without coupling the test to a random-init model's accuracy. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Addressed @cjluo-nv's review feedbackTwo rounds of comments handled — simplifications first, then the follow-up correctness/test nits. Summary of the changes (all in the latest commits):
|
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review at head (post-cd9eff8/latest). Verified against the files in the repo: this round closed the last substantive engineering gaps, so what's left is owner judgment plus one docs/design ask that has now gone three rounds without a change. No licensing changes; no prompt-injection attempts in the PR content.
Confirmed resolved this round (checked in-tree, not just from the diff):
tools/launcher/tests/test_core_extended.pynow hastest_reqs_inline_barrier/test_reqs_script_wrapped_inlineinTestRunJobsExtended, pinning the generatedrun.Script(inline=...)string:shlexquoting ('transformers<5'), theSLURM_LOCALIDbranch, the per-job/step/node marker, absence ofrm -f, and thebash <script> <args>wrap. This was the main outstanding test gap.- Barrier marker now folds
${SLURM_JOB_ID:-0}_${SLURM_STEP_ID:-0}_${SLURM_NODEID:-0}, so it's genuinely per-node/per-step andrm -f(and its race window) is gone; the shell reads correctly — rank 0's failed install short-circuits&&and the compoundifreturns non-zero, and waiting ranks fail on the trailing[ -f {marker} ].docs/configuration.mddocuments multi-nodereqs. lm_eval_hf.pygate is the simplified form: inlined bound pop and task parsing (with the empty---tasksregression fixed via[t.strip() for t in raw if t and t.strip()]), hoisted non-HF quant/sparsity validation reporting all offending flags, single fail-fast--output_pathcheck,mtime-selected results file,is_hfcovering all four HFLM aliases, and a comment justifying the disk read sinceHarnessCLI.execute() -> None.score_lower_boundintest_prune_minitron.pyis now0.0(no accuracy coupling);TestGetDefaultEnvenv isolation added;script/inlineexclusivity andargs+inlinerejected inrun_jobs(typed_target_tasks setself.scriptin__post_init__, so the new exactly-one check doesn't regress them).
Why this is a nudge rather than an approve:
- 💬 Design rationale (raised in three prior rounds; author's reply so far is the inline "
MLM_MODEL_CFG/QUANT_CFGare wrapper-only, N/A here"). That explains why the env vars are absent, not why the wrapper-free route was chosen: the PR body still doesn't compareinline/reqsagainst the existingcommon/**/*.shconvention (common/megatron_bridge/import/import.sh,common/megatron_lm/quantize/quantize.shalready cover this flow and could do the pip install + bash barrier) or typed_target_tasks. I re-readtools/launcher/CLAUDE.mdat head: untouched, "Key conventions" still says only "Scripts go incommon/" and "Adding a New Model Config" still says setMLM_MODEL_CFG/QUANT_CFG— exactly the rule these examples intentionally don't follow (the mismatch is what produced the two withdrawnMLM_MODEL_CFGfindings).docs/architecture.md/README.mdstill present typed tasks as the structured path. Contributors now have three ways to express a task command with no guidance; one or two sentences in the PR body plus a line inCLAUDE.mdcloses this. - 💬 Author consciously accepted the strict-checkpoint-key break from adding
"best"toMCoreMinitronSearcher.default_state_dict(pruning is one-shot). Reasonable, and I confirmedself.best = asdict(best)keepsBaseSearcher.search() -> SearchStateDictintact and is re-assigned on resume (the candidate loop still runs with restored scores), so the gate works after a resume. But resuming a pre-upgrademodelopt_pruning_scores/--prune_intermediate_ckptdir still dies onassert state_dict.keys() == self.state_dict().keys(), "Keys in checkpoint don't match!"with no hint to delete the dir — worth an explicit owner yes/no since the mitigation is one line. Also notesave_search_checkpoint()is never called after theself.bestassignment, so an on-disk completed checkpoint records"best": {}; harmless for the current gate (which reads the in-memory return), but a latent surprise for anyone reading the saved scores file. - Minor: the negative-path gates (
--accuracy_lower_bound/--score_lower_boundabove the score) remain untested; the author's argument that the passing path plus a trivialsys.exit(1)suffices is defensible, and CodeRabbit withdrew it — flagging only so the owner can agree consciously.
What does this PR do?
Type of change: new example + small launcher / modelopt-example features (backward compatible)
Adds end-to-end ModelOpt launcher pipelines for the Megatron-Bridge flow on Nemotron-3-Nano-30B-A3B, the minimal launcher features to run them wrapper-free from YAML, and an in-step accuracy gate for Minitron pruning.
New launcher examples (
tools/launcher/examples/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/):mbridge_prune.yaml— Minitron prune with an in-step MMLU gate → vLLM sanity gen (2 tasks)mbridge_quantize.yaml— FP8 quantize → unified-HF export → MMLU gate on the vLLM backend, which doubles as the deploy sanity check (3 tasks). Matches the tutorial.Prune accuracy gate (reuses the search's own score — no separate eval step):
modelopt/torch/prune/plugins/mcore_minitron.py—MCoreMinitronSearchernow stores the exported bestCandidateSubnetunderstate_dict["best"](additive; sits beside the existingsorted_layerskey).examples/megatron_bridge/prune_minitron.py— new--score_lower_bound: readspruning_scores["best"].scoreand exits non-zero if the exported model is below the floor. Score-agnostic (any--prune_score_func); rejected with--prune_export_config(manual pruning has no score).Launcher (
tools/launcher) — run single-node Megatron-Bridge one-liners directly from YAML:SandboxTask.inline— a command in the YAML, nocommon/**/*.shwrapper (single-line; folded scalar)SandboxTask.reqs/reqs_file— pip-install deps in the container before the command (shell-safe; on Slurm the install is rank-0-guarded so multi-rank tasks don't race)SlurmConfig.docker_user— local-Docker user (e.g.root); ignored on Slurmget_default_envhonorsHF_HOME/TRITON_CACHE_DIRenv overrides, so a non-CI user can point caches at a writable path (the shared/cicd/hf-cacheis owned by the CI account)argstogether withinlineexamples/llm_eval/lm_eval_hf.py:--accuracy_lower_bound— gate on the single requested task'sacc(used by the quantize MMLU step; exits non-zero if below)hfbackends, so--model vllmworks on a deployable quantized checkpointUsage
Testing
inline,reqs/reqs_file,docker_user, theargs+inlineguard, and example-resolve; ruff / mypy / bandit clean.tests/examples/megatron_bridge/test_prune_minitron.pynow passes--score_lower_bound=0.01to exercise the gate path on the tiny models.[score_gate] mmlu_10pct = 0.5196 >= 0.45 PASS; vLLM gen coherent.Detected ModelOpt fp8 checkpoint) → MMLU on the vLLM backendacc = 0.7077 >= 0.60 PASS.nemo:26.06through the same flow.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
transformers<5, sombridge_prune.yamlruns onnemo:26.04(26.06 drops it); quantize/export run onnemo:26.06.docker_user: rootis set on all example tasks — local-Docker only (ignored on Slurm), needed so downstream tasks can read task_0's root-owned checkpoints and to read the image's root-only/opt/Megatron-Bridge.enforce_eager=Trueto vLLM — for a run-once eval this skips ~17 min of CUDA-graph capture /torch.compilewith no accuracy change.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Bug Fixes